package com.dcits.agent.utils;

import net.bytebuddy.ByteBuddy;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.dynamic.loading.ClassReloadingStrategy;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.ElementMatchers;

import java.lang.reflect.Method;

/**
 * 运行时重定义class的工具类.
 *
 * @author caoyangjie
 */
public class RedefineClassUtils {

	private RedefineClassUtils() {
	}

	/**
	 * advice的class和原class进行混合，达到定制原class的目的.
	 *
	 * advice的class命名规则为：原class的名字 + 后缀{@code Advice}，这里的设计参考了mapstruct.
	 */
	public static Class<?> resolveAdviceClass(Class<?> sourceClass) {
		try {
			return Class.forName("advice."+sourceClass.getName() + "Advice");
		} catch (ClassNotFoundException e) {
			throw new IllegalStateException("source class's advice class not found", e);
		}
	}

	/**
	 * @throws IllegalArgumentException if there is no method match in source class
	 */
	public static void redefineMethodWithAdvice(Class<?> sourceClass, String methodName) {
		if (null == methodName) {
			throw new IllegalArgumentException("method name is null");
		}
		if (!existsMethodMatch(sourceClass, methodName)) {
			throw new IllegalArgumentException("there is no method '" + methodName + "' in " + sourceClass);
		}
		redefineMethodWithAdvice(sourceClass, ElementMatchers.named(methodName));
	}

	public static void redefineMethodWithAdvice(Class<?> sourceClass, ElementMatcher<? super MethodDescription> matcher) {
		final Class<?> adviceClass = resolveAdviceClass(sourceClass);
		new ByteBuddy().redefine(sourceClass)
				.visit(Advice.to(adviceClass).on(matcher))
				.make()
				.load(sourceClass.getClassLoader(), ClassReloadingStrategy.fromInstalledAgent())
				.getLoaded();
	}

	private static boolean existsMethodMatch(Class<?> clazz, String methodName) {
		Method[] methods = clazz.getMethods();
		for (Method method : methods) {
			if (method.getName().equals(methodName)) {
				return true;
			}
		}
		Method[] declaredMethods = clazz.getDeclaredMethods();
		for (Method method : declaredMethods) {
			if (method.getName().equals(methodName)) {
				return true;
			}
		}
		return false;
	}
}
